#include <bsl_params.h>
#include <bsl_errno.h>
#include <bsl_err_internal.h>
#include <crypt_algid.h>
#include <crypt_errno.h>
#include <crypt_types.h>
#include <securec.h>

#include "modes.h"
#include "err.h"

#include "wd.h"
#include "wd_cipher.h"
#include "wd_bmm.h"

void KAE_Mode_Free(KAE_Mode_Ctx *ctx)
{
    if (ctx == NULL) {
        return ;
    }

    if (ctx->pool != NULL && ctx->opData.in != NULL) {
        wd_free_blk(ctx->pool, ctx->opData.in);
    }

    if (ctx->pool != NULL && ctx->opData.out != NULL) {
        wd_free_blk(ctx->pool, ctx->opData.out);
    }

    if (ctx->pool != NULL && ctx->opData.iv != NULL) {
        wd_free_blk(ctx->pool, ctx->opData.iv);
    }

    if (ctx->cipherCtx != NULL) {
        wcrypto_del_cipher_ctx(ctx->cipherCtx);
    }

    if (ctx->pool != NULL) {
        wd_blkpool_destroy(ctx->pool);
    }

    wd_release_queue(&ctx->queue);

    free(ctx);
}

KAE_Mode_Ctx *KAE_Mode_NewCtx(int32_t algId)
{
    KAE_Mode_Ctx* ctx = (KAE_Mode_Ctx*) malloc(sizeof(KAE_Mode_Ctx));
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return NULL;
    }
    (void) memset_s(ctx, sizeof(KAE_Mode_Ctx), 0, sizeof(KAE_Mode_Ctx));

    ctx->queue.capa.alg = "cipher";
    wd_request_queue(&(ctx->queue));

    struct wd_blkpool_setup pool_setup;
    (void) memset_s(&pool_setup, sizeof(pool_setup), 0, sizeof(pool_setup));
    pool_setup.block_size = CIPHER_BLOCK_SIZE;
    pool_setup.block_num  = CIPHER_BLOCK_NUM;
    pool_setup.align_size = 64;
    ctx->pool = wd_blkpool_create(&ctx->queue, &pool_setup);
    if (ctx->pool == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }

    switch (algId) {
        case CRYPT_CIPHER_SM4_ECB:
        case CRYPT_CIPHER_SM4_CBC:
        case CRYPT_CIPHER_SM4_CTR:
        case CRYPT_CIPHER_SM4_XTS:
        case CRYPT_CIPHER_SM4_OFB: {
            ctx->setup.alg = WCRYPTO_CIPHER_SM4;
            ctx->cipherSize = SM4_BLOCK_SIZE;
            break;
        }

        case CRYPT_CIPHER_AES128_ECB:
        case CRYPT_CIPHER_AES128_CBC:
        case CRYPT_CIPHER_AES128_CTR:
        case CRYPT_CIPHER_AES128_OFB: {
            ctx->setup.alg = WCRYPTO_CIPHER_AES;
            ctx->cipherSize = AES128_BLOCK_SIZE;
            break;
        }

        case CRYPT_CIPHER_AES192_ECB:
        case CRYPT_CIPHER_AES192_CBC:
        case CRYPT_CIPHER_AES192_CTR:
        case CRYPT_CIPHER_AES192_OFB: {
            ctx->setup.alg = WCRYPTO_CIPHER_AES;
            ctx->cipherSize = AES192_BLOCK_SIZE;
            break;
        }

        case CRYPT_CIPHER_AES256_ECB:
        case CRYPT_CIPHER_AES256_CBC:
        case CRYPT_CIPHER_AES256_CTR:
        case CRYPT_CIPHER_AES256_OFB: {
            ctx->setup.alg = WCRYPTO_CIPHER_AES;
            ctx->cipherSize = AES256_BLOCK_SIZE;
            break;
        }

        default:
        {
            BSL_ERR_PUSH_ERROR(KAE_WD_NOT_SUPPORT);
            goto ERR;
        }
    }

    switch (algId) {
        case CRYPT_CIPHER_SM4_ECB:
        case CRYPT_CIPHER_AES128_ECB:
        case CRYPT_CIPHER_AES192_ECB:
        case CRYPT_CIPHER_AES256_ECB: {
            ctx->setup.mode = WCRYPTO_CIPHER_ECB;
            break;
        }

        case CRYPT_CIPHER_SM4_CBC:
        case CRYPT_CIPHER_AES128_CBC:
        case CRYPT_CIPHER_AES192_CBC:
        case CRYPT_CIPHER_AES256_CBC: {
            ctx->setup.mode = WCRYPTO_CIPHER_CBC;
            break;
        }

        case CRYPT_CIPHER_SM4_CTR:
        case CRYPT_CIPHER_AES128_CTR:
        case CRYPT_CIPHER_AES192_CTR:
        case CRYPT_CIPHER_AES256_CTR: {
            ctx->setup.mode = WCRYPTO_CIPHER_CTR;
            break;
        }

        case CRYPT_CIPHER_SM4_OFB:
        case CRYPT_CIPHER_AES128_OFB:
        case CRYPT_CIPHER_AES192_OFB:
        case CRYPT_CIPHER_AES256_OFB: {
            ctx->setup.mode = WCRYPTO_CIPHER_OFB;
            break;
        }

        case CRYPT_CIPHER_SM4_XTS: {
            ctx->setup.mode = WCRYPTO_CIPHER_XTS;
            break;
        }

        default:
        {
            BSL_ERR_PUSH_ERROR(KAE_WD_NOT_SUPPORT);
            goto ERR;
        }
    }

    ctx->setup.br.alloc = (void*) wd_alloc_blk;
    ctx->setup.br.free = (void*) wd_free_blk;
    ctx->setup.br.iova_map = (void*) wd_blk_iova_map;
    ctx->setup.br.iova_unmap = (void*) wd_blk_iova_unmap;
    ctx->setup.br.usr = ctx->pool;
    ctx->cipherCtx = wcrypto_create_cipher_ctx(&ctx->queue, &ctx->setup);
    if (ctx->cipherCtx == NULL) {
        BSL_ERR_PUSH_ERROR(KAE_WD_CREATE_CTX_FAIL);
        goto ERR;
    }
    ctx->pad = CRYPT_PADDING_PKCS7;

    return ctx;
ERR:
    KAE_Mode_Free(ctx);
    return NULL;
}

int32_t KAE_Mode_InitCtx(KAE_Mode_Ctx *ctx, const uint8_t *key, uint32_t keyLen, const uint8_t *iv, uint32_t ivLen, const BSL_Param *param, bool enc)
{
    if (ctx == NULL || ctx->cipherCtx == NULL || ctx->pool == NULL || key == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    (void) param;

    if (keyLen != ctx->cipherSize) {
        BSL_ERR_PUSH_ERROR(CRYPT_INVALID_ARG);
        return CRYPT_INVALID_ARG;
    }

    int32_t ret = wcrypto_set_cipher_key(ctx->cipherCtx, (uint8_t *) key, keyLen);
    if (ret != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_DO_CIPHER_FAIL);
        return KAE_WD_DO_CIPHER_FAIL;
    }

    ctx->opData.op_type = enc ? WCRYPTO_CIPHER_ENCRYPTION : WCRYPTO_CIPHER_DECRYPTION;

    ctx->opData.in = wd_alloc_blk(ctx->pool);
    if (ctx->opData.in == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }

    ctx->opData.out = wd_alloc_blk(ctx->pool);
    if (ctx->opData.out == NULL) {
        wd_free_blk(ctx->pool ,ctx->opData.in);
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        return CRYPT_MEM_ALLOC_FAIL;
    }

    ctx->opData.iv = wd_alloc_blk(ctx->pool);
    if (ctx->opData.iv == NULL) {
        ret = CRYPT_MEM_ALLOC_FAIL;
        BSL_ERR_PUSH_ERROR(CRYPT_MEM_ALLOC_FAIL);
        goto ERR;
    }
    ctx->opData.iv_bytes = 0;

    switch(ctx->setup.mode) {
        case WCRYPTO_CIPHER_CBC:
	    case WCRYPTO_CIPHER_CTR:
	    case WCRYPTO_CIPHER_XTS:
	    case WCRYPTO_CIPHER_OFB: {
            if (iv == NULL) {
                ret = CRYPT_NULL_INPUT;
                BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
                goto ERR;
            }

            if (ivLen != ctx->cipherSize) {
                ret = CRYPT_INVALID_ARG;
                BSL_ERR_PUSH_ERROR(CRYPT_INVALID_ARG);
                goto ERR;
            }
            (void) memcpy_s(ctx->opData.iv, ivLen, iv, ivLen);
            ctx->opData.iv_bytes = ivLen;
            break;
        }
        default: break;
    }
    return CRYPT_SUCCESS;
ERR:
    wd_free_blk(ctx->pool ,ctx->opData.in);
    wd_free_blk(ctx->pool ,ctx->opData.out);
    return ret;
}

int32_t KAE_Mode_Update(KAE_Mode_Ctx *ctx, const uint8_t *in, uint32_t inLen, uint8_t *out, uint32_t *outLen)
{
    if ((ctx == NULL) || (ctx->cipherCtx == NULL) || (in == NULL && inLen != 0)) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (inLen == 0) {
        return CRYPT_SUCCESS;
    }

    (*outLen) = 0;

    int cnt = 0;

    if (ctx->remain > 0)
    {
        if (inLen + ctx->remain > CIPHER_CACHE_SIZE)
        {
            (void) memcpy_s(ctx->block + ctx->remain, CIPHER_CACHE_SIZE - ctx->remain, in, CIPHER_CACHE_SIZE - ctx->remain);
            in += CIPHER_CACHE_SIZE - ctx->remain;
            inLen -= (CIPHER_CACHE_SIZE - ctx->remain);
            ctx->remain = 0;
            ctx->opData.in_bytes = CIPHER_CACHE_SIZE;
            (void) memcpy_s(ctx->opData.in, CIPHER_CACHE_SIZE, ctx->block, ctx->opData.in_bytes);
            if (wcrypto_do_cipher(ctx->cipherCtx, &ctx->opData, NULL)  != CRYPT_SUCCESS) {
                BSL_ERR_PUSH_ERROR(KAE_WD_DO_CIPHER_FAIL);
                return KAE_WD_DO_CIPHER_FAIL;
            }
            (void) memcpy_s(out, CIPHER_CACHE_SIZE, ctx->opData.out, CIPHER_CACHE_SIZE);
            out += CIPHER_CACHE_SIZE;
            (*outLen) = CIPHER_CACHE_SIZE;
        }
        else
        {
            memcpy(ctx->block + ctx->remain, in, inLen);
            ctx->remain += inLen;
            return CRYPT_SUCCESS;
        }
    }

    while (inLen - cnt > CIPHER_CACHE_SIZE)
    {
        ctx->opData.in_bytes = CIPHER_CACHE_SIZE;
        (void) memcpy_s(ctx->opData.in, CIPHER_CACHE_SIZE, ctx->block, ctx->opData.in_bytes);
        if (wcrypto_do_cipher(ctx->cipherCtx, &ctx->opData, NULL)  != CRYPT_SUCCESS) {
            BSL_ERR_PUSH_ERROR(KAE_WD_DO_CIPHER_FAIL);
            return KAE_WD_DO_CIPHER_FAIL;
        }
        (void) memcpy_s(out, CIPHER_CACHE_SIZE, ctx->opData.out, CIPHER_CACHE_SIZE);
        out += CIPHER_CACHE_SIZE;
        (*outLen) += CIPHER_CACHE_SIZE;
        in += CIPHER_CACHE_SIZE;
        cnt += CIPHER_CACHE_SIZE;
    }

    ctx->remain = inLen - cnt;
    memcpy(ctx->block, in, ctx->remain);

    return CRYPT_SUCCESS;
}

int32_t KAE_Mode_Final_PAD(KAE_Mode_Ctx *ctx, uint8_t *out, uint32_t *outLen)
{
    uint32_t paddingSize = 0;
    if (ctx->opData.op_type == WCRYPTO_CIPHER_ENCRYPTION) {
        paddingSize = (ctx->cipherSize - (ctx->remain % ctx->cipherSize)) % ctx->cipherSize;
        (void) memset_s(ctx->block + ctx->remain, paddingSize, paddingSize, paddingSize);
    }

    ctx->opData.in_bytes = ctx->remain + paddingSize;

    if (*outLen < ctx->opData.in_bytes) {
        BSL_ERR_PUSH_ERROR(CRYPT_MODE_BUFF_LEN_NOT_ENOUGH);
        return CRYPT_MODE_BUFF_LEN_NOT_ENOUGH;
    }

    memcpy(ctx->opData.in, ctx->block, ctx->opData.in_bytes);
    if (wcrypto_do_cipher(ctx->cipherCtx, &ctx->opData, NULL)  != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_DO_CIPHER_FAIL);
        return KAE_WD_DO_CIPHER_FAIL;
    }

    if (ctx->opData.op_type == WCRYPTO_CIPHER_ENCRYPTION) {
        (void) memcpy_s(out, CIPHER_CACHE_SIZE, ctx->opData.out, ctx->opData.in_bytes);
        (*outLen) = ctx->opData.in_bytes;
    }

    if (ctx->opData.op_type ==  WCRYPTO_CIPHER_DECRYPTION) {
        uint8_t *opOut = (uint8_t *) ctx->opData.out;
        paddingSize = *(opOut + ctx->opData.in_bytes - 1);
        (void) memcpy_s(out, CIPHER_CACHE_SIZE, ctx->opData.out, ctx->opData.in_bytes - paddingSize);
        (*outLen) = ctx->opData.in_bytes - paddingSize;
    }

    return CRYPT_SUCCESS;
}

int32_t KAE_Mode_Final_NO_PAD(KAE_Mode_Ctx *ctx, uint8_t *out, uint32_t *outLen)
{
    ctx->opData.in_bytes = ctx->remain;
    memcpy(ctx->opData.in, ctx->block, ctx->opData.in_bytes);
    if (wcrypto_do_cipher(ctx->cipherCtx, &ctx->opData, NULL)  != CRYPT_SUCCESS) {
        BSL_ERR_PUSH_ERROR(KAE_WD_DO_CIPHER_FAIL);
        return KAE_WD_DO_CIPHER_FAIL;
    }
    memcpy(out, ctx->opData.out, ctx->opData.in_bytes);
    (*outLen) = ctx->opData.in_bytes;
    return CRYPT_SUCCESS;
}

int32_t KAE_Mode_Final(KAE_Mode_Ctx *ctx, uint8_t *out, uint32_t *outLen)
{
    if ((ctx == NULL) || (ctx->cipherCtx == NULL) || (out == NULL && outLen != 0 )) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (outLen == 0) {
        return CRYPT_SUCCESS;
    }

    if (ctx->setup.mode == WCRYPTO_CIPHER_CTR || ctx->setup.mode == WCRYPTO_CIPHER_OFB) {
        return KAE_Mode_Final_NO_PAD(ctx, out, outLen);
    } else {
        return KAE_Mode_Final_PAD(ctx, out, outLen);
    }
}

int32_t MODES_SetPaddingCheck(int32_t pad)
{
    if (pad >= CRYPT_PADDING_MAX_COUNT || pad < CRYPT_PADDING_NONE) {
        BSL_ERR_PUSH_ERROR(CRYPT_MODES_PADDING_NOT_SUPPORT);
        return CRYPT_MODES_PADDING_NOT_SUPPORT;
    }
    return CRYPT_SUCCESS;
}

int32_t KAE_Mode_SetIv(KAE_Mode_Ctx *ctx, void *val, uint32_t valLen)
{
    if (val == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    if (valLen != ctx->cipherSize) {
        BSL_ERR_PUSH_ERROR(CRYPT_MODES_IVLEN_ERROR);
        return CRYPT_MODES_IVLEN_ERROR;
    }

    if (memcpy_s(ctx->opData.iv, valLen, val, valLen) != EOK) {
        BSL_ERR_PUSH_ERROR(CRYPT_SECUREC_FAIL);
        return CRYPT_SECUREC_FAIL;
    }

    ctx->opData.iv_bytes = valLen;
    return CRYPT_SUCCESS;
}

int32_t KAE_Mode_GetIv(KAE_Mode_Ctx *ctx, void *val, uint32_t valLen)
{
    if (val == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    uint32_t ivLen = ctx->cipherSize;

    if (valLen != ivLen) {
        BSL_ERR_PUSH_ERROR(CRYPT_MODE_ERR_INPUT_LEN);
        return CRYPT_MODE_ERR_INPUT_LEN;
    }

    if (memcpy_s(val, valLen, ctx->opData.iv, valLen) != EOK) {
        BSL_ERR_PUSH_ERROR(CRYPT_SECUREC_FAIL);
        return CRYPT_SECUREC_FAIL;
    }

    return CRYPT_SUCCESS;
}

int32_t KAE_Mode_Ctrl(KAE_Mode_Ctx *ctx, int32_t cmd, void *val, uint32_t valLen)
{
    int32_t ret;
    if (ctx == NULL || ctx->cipherCtx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    switch (cmd) {
        case CRYPT_CTRL_SET_PADDING:
            ret = MODES_SetPaddingCheck(*(int32_t *)val);
            if (ret != CRYPT_SUCCESS) {
                return ret;
            }
            ctx->pad = *(int32_t *) val;
            return CRYPT_SUCCESS;
        case CRYPT_CTRL_GET_PADDING:
            *(int32_t *)val = ctx->pad;
            return CRYPT_SUCCESS;
        case CRYPT_CTRL_GET_BLOCKSIZE:
            if (val == NULL || valLen != sizeof(uint32_t)) {
                BSL_ERR_PUSH_ERROR(CRYPT_MODE_ERR_INPUT_LEN);
                return CRYPT_MODE_ERR_INPUT_LEN;
            }
            *(int32_t *)val =ctx->cipherSize;
            return CRYPT_SUCCESS;
        case CRYPT_CTRL_REINIT_STATUS:
            return KAE_Mode_SetIv(ctx, val, valLen);
        case CRYPT_CTRL_GET_IV:
            return KAE_Mode_GetIv(ctx, (uint8_t *) val, valLen);
        default:
            BSL_ERR_PUSH_ERROR(CRYPT_MODES_CTRL_TYPE_ERROR);
            return CRYPT_MODES_CTRL_TYPE_ERROR;
    }
}

int32_t KAE_Mode_Deinit(KAE_Mode_Ctx *ctx)
{
    if (ctx == NULL) {
        BSL_ERR_PUSH_ERROR(CRYPT_NULL_INPUT);
        return CRYPT_NULL_INPUT;
    }

    wd_free_blk(ctx->pool, ctx->opData.in);
    wd_free_blk(ctx->pool, ctx->opData.out);
    wd_free_blk(ctx->pool, ctx->opData.iv);
    (void) memset_s(&ctx->opData, sizeof(struct wcrypto_cipher_op_data), 0, sizeof(struct wcrypto_cipher_op_data));

    wcrypto_del_cipher_ctx(ctx->cipherCtx);
    ctx->cipherCtx = wcrypto_create_cipher_ctx(&ctx->queue, &ctx->setup);
    if (ctx->cipherCtx == NULL) {
        BSL_ERR_PUSH_ERROR(KAE_WD_CREATE_CTX_FAIL);
        return KAE_WD_CREATE_CTX_FAIL;
    }
    return CRYPT_SUCCESS;
}